import os
import cv2
from pathlib import Path

# 原始数据集路径
existing_labels_dir = Path(r"D:\Datasets\dataset2CheckNew\labels")
existing_images_dir = Path(r"D:\Datasets\dataset2CheckNew\images")

# 新分类数据集路径
classifier_dataset_dir = Path(r"D:\Datasets\dataset2CheckNew\classifier_dataset")
hel_dir = classifier_dataset_dir / "hel"
nohel_dir = classifier_dataset_dir / "nohel"
hel_dir.mkdir(exist_ok=True, parents=True)
nohel_dir.mkdir(exist_ok=True, parents=True)

# 支持的图像扩展名（可根据需要添加）
SUPPORTED_IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}

# 类别定义
CLASS_MAP = {
    0: hel_dir,      # hel
    1: nohel_dir     # nohel
}

counter = 0
missing_images = 0
failed_reads = 0
invalid_crops = 0

print("开始处理标签和图像文件...")

for label_file in existing_labels_dir.glob("*.txt"):
    image_stem = label_file.stem

    # 尝试所有支持的扩展名
    image_path = None
    for ext in SUPPORTED_IMAGE_EXTENSIONS:
        potential_path = existing_images_dir / f"{image_stem}{ext}"
        if potential_path.exists():
            image_path = potential_path
            break

    if image_path is None:
        print(f"警告：找不到图像 {image_stem} 的任何支持格式")
        missing_images += 1
        continue

    # 读取图像
    image = cv2.imread(str(image_path))
    if image is None:
        print(f"警告：无法读取图像 {image_path}")
        failed_reads += 1
        continue

    height, width = image.shape[:2]

    # 读取YOLO标注
    try:
        with open(label_file, 'r', encoding='utf-8') as f:
            lines = f.readlines()
    except Exception as e:
        print(f"错误：无法读取标签文件 {label_file} - {e}")
        continue

    for line in lines:
        parts = line.strip().split()
        if not parts:
            continue

        try:
            class_id = int(parts[0])
        except ValueError:
            print(f"警告：无效的类别ID格式 {parts[0]} in {label_file}")
            continue

        if class_id not in CLASS_MAP:
            continue  # 忽略不关心的类别

        # 解析YOLO格式 (x_center, y_center, w, h) - 归一化
        try:
            x_center_norm = float(parts[1])
            y_center_norm = float(parts[2])
            w_norm = float(parts[3])
            h_norm = float(parts[4])
        except (ValueError, IndexError):
            print(f"警告：无效的标注格式 in {label_file}: {line}")
            continue

        # 转换为像素坐标
        x_center = int(x_center_norm * width)
        y_center = int(y_center_norm * height)
        bbox_width = int(w_norm * width)
        bbox_height = int(h_norm * height)

        # 计算边界框
        x1 = max(0, x_center - bbox_width // 2)
        y1 = max(0, y_center - bbox_height // 2)
        x2 = min(width, x1 + bbox_width)
        y2 = min(height, y1 + bbox_height)

        if x1 >= x2 or y1 >= y2:
            print(f"警告：无效的裁剪区域 [{x1},{y1},{x2},{y2}] in {image_path}")
            invalid_crops += 1
            continue

        # 裁剪图像
        cropped = image[y1:y2, x1:x2]
        if cropped.size == 0:
            invalid_crops += 1
            continue

        # 保存路径
        save_dir = CLASS_MAP[class_id]
        save_filename = f"{image_stem}_crop_{counter:04d}.jpg"  # 统一保存为 .jpg
        save_path = save_dir / save_filename

        success = cv2.imwrite(str(save_path), cropped)
        if not success:
            print(f"警告：无法保存图像 {save_path}")
        else:
            counter += 1

# 汇总输出
print("\n✅ 处理完成！")
print(f"共裁剪并保存了 {counter} 个图像。")
if missing_images:
    print(f"⚠️  找不到图像文件：{missing_images} 个")
if failed_reads:
    print(f"⚠️  图像读取失败：{failed_reads} 个")
if invalid_crops:
    print(f"⚠️  无效或跳过的裁剪区域：{invalid_crops} 个")